feat(accuracy): restore codegen grade concurrency (AIP-1094) - #1237

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

tests/unit/accuracy/test_codegen_worker_client.py (1)

82-95: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

TestSerialization no longer tests serialization; consider moving this test into TestConcurrency.

The comment at Line 83-84 states this class previously verified serialized execution and now verifies the opposite (concurrent execution without a lock). Keeping the old class name is confusing since a TestConcurrency class already exists immediately below with the same concurrency scope. Move test_concurrent_grades_return_correct_results into TestConcurrency and drop the now-empty TestSerialization class, or rename the class to reflect its new purpose.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unit/accuracy/test_codegen_worker_client.py` around lines 82 - 95, Move
test_concurrent_grades_return_correct_results from TestSerialization into the
existing TestConcurrency class, preserving its assertions and cleanup. Remove
the now-empty TestSerialization class and update any obsolete serialization
wording associated with the test.

src/aiperf/accuracy/graders/_codegen_worker.py (1)

80-158: 🚀 Performance & Scalability | 🔵 Trivial

handle_batch correctly isolates per-request failures and preserves order.

The malformed-JSON, non-dict, missing-field, batch-exception, and per-item-metric-exception paths all populate responses at the original index and are covered by the referenced unit tests (test_malformed_request_in_batch_does_not_affect_others, test_batch_exception_returns_error_for_all, test_response_order_matches_request_order). The blind except Exception at Line 134 and Line 151 flagged by Ruff is intentional per the docstring's "Never raises" guarantee, so isolated per-request/per-batch failures do not take down the worker loop.

One scalability note: _drain_buffered (called from run_worker_loop) can grow reqs unboundedly while queued input exists, and no response is written until the entire batch's codegen_fn call completes. Under a burst of many concurrent grade_codegen calls, this can inflate per-request latency and peak memory for large evaluation_sample/generated_code payloads. Consider capping batch size if this becomes an issue in practice.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/aiperf/accuracy/graders/_codegen_worker.py` around lines 80 - 158,
Consider adding a maximum batch-size limit to the _drain_buffered flow used by
run_worker_loop so reqs cannot grow without bound while input remains queued.
Process buffered requests in capped batches and preserve response ordering and
existing handle_batch error isolation.

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

src/aiperf/accuracy/graders/_codegen_worker.py (1)

191-196: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add a pipe-backed partial-line test.

Use io.BufferedReader over os.pipe(). Write one JSONL request in two writes and assert that the request completes after the newline arrives. Existing BytesIO tests do not exercise peek(0) or the partial-line guard.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/aiperf/accuracy/graders/_codegen_worker.py` around lines 191 - 196, Add a
pipe-backed partial-line test for the worker’s JSONL request-reading path, using
io.BufferedReader over os.pipe(). Split one request across two writes, verify
the first partial write does not complete processing, then write the newline and
assert the request completes successfully; keep the test focused on exercising
peek(0) and the partial-line guard around the worker read loop.

Source: Learnings

tests/unit/accuracy/test_codegen_worker.py (4)

135-147: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Reuse the existing batch fixture instead of redefining it.

_nan_inf is identical to _fake_codegen_batch_ok at lines 32-38. Pass _fake_codegen_batch_ok and keep only the custom _nan_compute.

♻️ Proposed simplification

     def test_non_finite_metric_values_are_dropped(self) -> None:
         # NaN/Inf must not cross the JSONL boundary (repo NaN/Inf discipline).
-        def _nan_inf(
-            samples: list, generations: list, **_kwargs: Any
-        ) -> tuple[dict[str, Any], dict[int, list]]:
-            n = len(samples)
-            return {"pass@1": 1.0}, {i: [[True]] for i in range(n)}
-
         def _nan_compute(results: dict, **_kwargs: Any) -> dict[str, Any]:
             return {"pass@1": float("nan"), "extra": float("inf"), "ok": 1.0}
 
         req = {"id": 9, "evaluation_sample": [{}], "generated_code": [["x"]]}
-        resps = worker.handle_batch([req], _nan_inf, _nan_compute)
+        resps = worker.handle_batch([req], _fake_codegen_batch_ok, _nan_compute)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unit/accuracy/test_codegen_worker.py` around lines 135 - 147, Update
test_non_finite_metric_values_are_dropped to remove the duplicate _nan_inf
helper and pass the existing _fake_codegen_batch_ok fixture to
worker.handle_batch, retaining only the custom _nan_compute behavior.

---

202-203: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Move threading and time imports to the file top.

The repository test guidelines require imports at the file top.

As per coding guidelines: "keep imports at the top".

♻️ Proposed change

     def test_partial_jsonl_line_is_deferred_to_next_cycle(self) -> None:
         # Exercises the O_NONBLOCK peek(0) + partial-line guard in _drain_buffered.
         # A partial write (no trailing newline) must not be submitted as a request;
         # only after the newline arrives should the line be processed.
-        import threading
-        import time
-
         req = self._req(42)

Add at the file top:

import threading
import time

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unit/accuracy/test_codegen_worker.py` around lines 202 - 203, Move the
threading and time imports from their local position near the affected test into
the module-level import section at the top of
tests/unit/accuracy/test_codegen_worker.py, leaving their usage unchanged.

Source: Coding guidelines

---

241-244: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Annotate the codegen_fn parameter.

codegen_fn has no type hint. TestRunWorkerLoopBatch._run already uses Callable[..., tuple[dict[str, Any], dict[int, list]]]. Use the same annotation here.

As per coding guidelines: "Add type hints to every function parameter and return value".

♻️ Proposed change

-    def _run(self, requests: list[bytes], codegen_fn) -> list[dict]:
+    def _run(
+        self,
+        requests: list[bytes],
+        codegen_fn: Callable[..., tuple[dict[str, Any], dict[int, list]]],
+    ) -> list[dict[str, Any]]:

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unit/accuracy/test_codegen_worker.py` around lines 241 - 244, Update
the _run method’s codegen_fn parameter annotation to match
TestRunWorkerLoopBatch._run: Callable[..., tuple[dict[str, Any], dict[int,
list]]]. Preserve the existing requests annotation and list[dict] return
annotation.

Source: Coding guidelines

---

210-237: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low value

Guard the shared BytesIO against cross-thread access.

The worker thread writes to out while the main thread calls out.seek(0) and out.read() at lines 224-225. io.BytesIO keeps one shared file position, so a concurrent write would land at the position set by the main thread. The current ordering is safe because the worker writes only after line 228, but the pattern breaks silently if the timing changes. Consider a small lock-protected writer wrapper, or snapshot with out.getvalue() instead of seek/read.

♻️ Proposed change

-            assert t.is_alive()
-            out.seek(0)
-            assert out.read() == b""  # nothing written yet
+            assert t.is_alive()
+            assert out.getvalue() == b""  # nothing written yet
@@
         assert not t.is_alive()
-        out.seek(0)
-        resps = [orjson.loads(ln) for ln in out if ln.strip()]
+        resps = [
+            orjson.loads(ln) for ln in out.getvalue().splitlines() if ln.strip()
+        ]

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unit/accuracy/test_codegen_worker.py` around lines 210 - 237, Protect
the shared BytesIO access in the worker-thread test around the output inspection
after starting worker.run_worker_loop. Replace the main thread’s seek/read
pattern with a position-independent snapshot such as getvalue(), and use that
snapshot for both the empty-output assertion and final response parsing so
concurrent writes cannot alter the shared file position.

.gitignore (1)

60-60: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Remove the duplicate ignore rule.

Line 56 already ignores .worktrees/.

♻️ Proposed change

 tests/scripts/.chaos_runs/
-.worktrees/

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In @.gitignore at line 60, Remove the duplicate `.worktrees/` entry from the
changed section, keeping the existing ignore rule at line 56 as the sole rule.

src/aiperf/accuracy/graders/_codegen_worker_client.py (2)

296-300: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Extract the pending-future failure loop.

The same loop appears in _handle_fault at lines 251-254. Extract a small helper so both paths stay consistent.

♻️ Proposed change

+    def _fail_pending(self, message: str) -> None:
+        for fut in list(self._pending.values()):
+            if not fut.done():
+                fut.set_exception(CodegenWorkerError(message))
+        self._pending.clear()
+
     async def aclose(self) -> None:
-        for fut in list(self._pending.values()):
-            if not fut.done():
-                fut.set_exception(CodegenWorkerError("grading worker closed"))
-        self._pending.clear()
+        self._fail_pending("grading worker closed")
         await self._kill()

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/aiperf/accuracy/graders/_codegen_worker_client.py` around lines 296 -
300, Extract the pending-future failure loop from the current close path into a
small helper near the relevant methods, then call that helper from both the
close flow and `_handle_fault`. Preserve the existing behavior of setting
CodegenWorkerError("grading worker closed") on unfinished futures and clearing
`_pending` exactly once.

---

218-240: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win

Fault the worker if the reader raises an unexpected exception.

The try block only handles ValueError, ConnectionError, BrokenPipeError, orjson.JSONDecodeError, and CancelledError. Any other exception ends the reader task silently. Every pending caller then waits for its full timeout, and the task exception is only reported when the task is garbage collected.

Add a final except Exception that calls _handle_fault(). Note that BrokenPipeError is a subclass of ConnectionError, so it is redundant in the tuple at line 222.

🛠️ Proposed change

         except asyncio.CancelledError:
             pass
+        except Exception:
+            # A reader crash would leave every pending caller waiting for its
+            # timeout, so convert it into a worker fault.
+            await self._handle_fault()

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/aiperf/accuracy/graders/_codegen_worker_client.py` around lines 218 -
240, Update the reader loop in the worker client to catch unexpected exceptions
with a final except Exception branch and await _handle_fault() before returning
or completing. Remove the redundant BrokenPipeError entry from the existing
ConnectionError tuple, while preserving the current CancelledError handling and
fault behavior for known read, parse, and dispatch failures.

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes